package sample.methodadvice;

import org.objectweb.asm.*;
import org.objectweb.asm.commons.AdviceAdapter;
import util.CustomClassLoader;
import util.FileUtils;

import java.io.IOException;
import java.util.Date;

import static org.objectweb.asm.Opcodes.ACC_ABSTRACT;
import static org.objectweb.asm.Opcodes.ACC_NATIVE;

public class MethodAroundDemo {
    public static void main(String[] args) throws IOException {
        byte[] bytes = invoke(HelloWorld.class);
        //  output
        FileUtils.writeBytes(Type.getInternalName(HelloWorld.class) + "MethodAround.class", bytes);
        //  replace
        FileUtils.writeBytes(HelloWorld.class.getResource("/").getPath() + Type.getInternalName(HelloWorld.class) + ".class", bytes);

        invokeMethod(bytes);
    }

    private static byte[] invoke(Class<?> clazz) throws IOException {
        ClassReader classReader = new ClassReader(clazz.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES);
        MethodAroundVisitor visitor = new MethodAroundVisitor(Opcodes.ASM9, classWriter);
        classReader.accept(visitor, ClassReader.EXPAND_FRAMES);
        byte[] bytes = classWriter.toByteArray();
        return bytes;
    }

    private static void invokeMethod(byte[] bytes) {
        CustomClassLoader classLoader = new CustomClassLoader();
        Class clazz = classLoader.defineClass(HelloWorld.class.getName(), bytes);
        Object instance = classLoader.invokeInstance(clazz, new Class[]{String.class, int.class}, new Object[]{"ydx", 24});
        classLoader.invokeMethod(clazz, instance, new Class[]{long.class, Object.class}, "test", new Object[]{80276415, new Date()});
    }

    private static class MethodAroundVisitor extends ClassVisitor {
        public MethodAroundVisitor(int api, ClassVisitor classVisitor) {
            super(api, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            if (mv != null) {
                boolean isAbstractMethod = (access & ACC_ABSTRACT) != 0;
                boolean isNativeMethod = (access & ACC_NATIVE) != 0;
                if (!isAbstractMethod && !isNativeMethod) {
                    mv = new MethodAroundAdapter(api, mv, access, name, descriptor);
                }
            }
            return mv;
        }

        private static class MethodAroundAdapter extends AdviceAdapter {
            protected MethodAroundAdapter(int api, MethodVisitor methodVisitor, int access, String name, String descriptor) {
                super(api, methodVisitor, access, name, descriptor);
            }

            @Override
            protected void onMethodEnter() {
                printMessage(String.format("Method Enter: %s %s", getName(), methodDesc));
                Type[] types = getArgumentTypes();
                int typesLength = types.length;
                for (int i = 0; i < typesLength; i++) {
                    Type type = types[i];
                    loadArg(i);
                    box(type);
                    printValueOnStack("(Ljava/lang/Object;)V");
                }
            }

            @Override
            protected void onMethodExit(int opcode) {
                printMessage(String.format("Method Exit: %s %s", getName(), methodDesc));
                if (opcode == ATHROW) {
                    super.visitLdcInsn("abnormal return");
                } else if (opcode == RETURN) {
                    super.visitLdcInsn("return void");
                } else if (opcode == ARETURN) {
                    dup();
                } else {
                    if (opcode == LRETURN || opcode == DRETURN) {
                        dup2();
                    } else {
                        dup();
                    }
                    box(getReturnType());
                }
                printValueOnStack("(Ljava/lang/Object;)V");
            }

            private void printMessage(String message) {
                super.visitLdcInsn(message);
                super.visitMethodInsn(INVOKESTATIC, "util/ParameterUtils", "printText", "(Ljava/lang/String;)V", false);
            }

            private void printValueOnStack(String descriptor) {
                super.visitMethodInsn(INVOKESTATIC, "util/ParameterUtils", "printValueOnStack", descriptor, false);
            }
        }
    }
}
